
import arrow
from bson import ObjectId
from app.db.models.asset_tasks import get_monitor_urls
from app.db.models.tasks import IgnoreWarning, ChangeCheckBasepage
from app.db.models.wass import ScanEvents
from app.db.models.asset_tasks import AssetTasks
from app.libs.enums import TaskType, ChangeCheckType, TaskTriggerType

from . import BaseHandler
from .utils import get_warning_id
from .events import Event
from ...libs.score import Score


class ChangCheckHandler(BaseHandler):

    def handle(self):
        return self._handle_change_check(task_settings=self.task_settings, response=self.response)

    def _handle_change_check(self, task_settings, response):
        statistics = {
            "urls": len(get_monitor_urls(task_session_id=self.task.taskSessionId, is_source_data=True)),
            "scan_count": get_monitor_urls(
                task_session_id=self.task.taskSessionId, is_scan_count=True, is_source_data=True),
            "is_update_baseline": response.get(
                "statistics", {}).get("change_check", {}).get("is_update_baseline", False)
        }
        result = {
            "addition": {
                "statistics": statistics,
                "site_info": {"waf_list": response.get("site_info", {}).get("waf_list", [])},
                "target_ip": response.get("target_ip", "")
            },
            "warnings": []
        }
        if target_status := Event._target_status(response.get("site_info", {})):
            result["targetStatus"] = target_status
        severity_map = Score.severityMap
        # 篡改监测忽略威胁逻辑与其他监测项不太一致， 提前处理
        ignore_warnings = [f"{_.detail.affects}||{_.detail.category}" for _ in IgnoreWarning.objects.filter(
            jobId=self.task.jobId)]
        for item in response.get("change_check", []):
            affects = item.get("url")
            warning = {
                "id": get_warning_id(
                    risk_title="change_check",
                    task_type=TaskType.change_check.value,
                    url=affects,
                    job_id=self.task.jobId
                ),
                "level": severity_map.get(item["level"], {}).get("level"),
                "levelName": severity_map.get(item["level"], {}).get("zh_CN"),
                "title": "篡改",
                "category": "change_check",
                "affects": affects,
                "foundAt": item.get("found_at"),
                "isResolved": False,
                "detail": {},
                "_id": ObjectId()
            }
            for change_check_type in ChangeCheckType:
                # 已经忽略的变更类型不再记录
                if f"{affects}||{change_check_type.value}" in ignore_warnings:
                    continue
                if item.get(f"{change_check_type.value}_result", {}).get("change_ratio", 0) > 0:
                    warning["detail"].update({
                        change_check_type.value: item.get(f"{change_check_type.value}_result", {})
                    })
            if not warning["detail"]:
                continue
            result["warnings"].append(warning)
        return result


# 篡改监测更新样本
def update_baseline(trigger_type, task_id, task_session_id):
    if not task_id:
        return
    old_pages = ChangeCheckBasepage.objects.filter(task_id=task_id)
    asset_task_obj = AssetTasks.objects.filter(changeCheckSettings__taskId=task_id).first()
    if not asset_task_obj:
        return
    old_page_maps = {f"{p.task_id}_{p.url}": p.to_mongo().to_dict() for p in old_pages}
    new_pages = ScanEvents._get_db()["scan_page_baselines"].find({"task_id": task_id})
    new_page_maps = {f"{p['task_id']}_{p['url']}": p for p in new_pages}
    if trigger_type == TaskTriggerType.manual.value:
        update_type = "manual"
    else:
        update_type = "auto"
    results = []
    for _, item in new_page_maps.items():
        results.append({
            "uid": asset_task_obj.uid,
            "job_id": asset_task_obj.jobId,
            "asset_task_id": asset_task_obj.id,
            "task_id": task_id,
            "url": item.get("url"),
            "status": "completed",
            "update_type": update_type,
            "task_session_id": task_session_id,
            "change_count": old_page_maps.get(_, {}).get("change_count", 0) + 1,
            "text": item.get("text", {}),
            "structure": item.get("structure", {}),
            "resources": item.get("resources", []),
            "created_at": old_page_maps.get(_, {}).get("created_at", arrow.utcnow().datetime),
            "updated_at": arrow.utcnow().datetime
        })
    ChangeCheckBasepage.objects(task_id=task_id).delete()
    ChangeCheckBasepage._get_collection().insert_many(results)
